package com.catmiao.spark.sql.test

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Encoder, Encoders, Row, SparkSession, functions}

/**
 * @title: Spark01_sql_test01
 * @projectName spark_study
 * @description: TODO
 * @author ChengMiao
 * @date 2024/3/20 10:44
 */
object Spark03_sql_UDAF2 {

  def main(args: Array[String]): Unit = {
    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("spark_sql_01_test")

    // 创建 Spark Session对象
    val sparkSession: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
    // 涉及转换操作，需要引入转换规则

    val df: DataFrame = sparkSession.read.json("datas/sql/user.json")


    df.createOrReplaceTempView("user")
    // 内置函数 拼接常量前缀
    //    sparkSession.sql("select CONCAT('Name:',username),age from user").show
    // 自定义函数 添加前缀 注册自定义函数
    sparkSession.udf.register("ageAvg", functions.udaf(new MyAvgUDAF))
    sparkSession.sql("select ageAvg(age) from user").show




    // 关闭环境
    sparkSession.close()
  }


  /**
   * 自定义聚合函数类 计算年龄的平均值
   * 1. 继承 Aggregator 定义泛型
   *      - IN : 输入的数据类型
   *      - BUF: 缓冲区
   *      - OUT: 输出的数据类型
   */

  case class Buff(var total: Long, var count: Long) // 缓冲区数据类型

  class MyAvgUDAF extends Aggregator[Long, Buff, Long] {
    /**
     * 初始值，缓冲区初始化
     *
     * @return
     */

    override def zero: Buff = {
      Buff(0, 0)
    }


    /**
     * 聚合，根据输入数据更新缓冲区数据
     */
    override def reduce(b: Buff, a: Long): Buff = {
      b.count = b.count + 1
      b.total = b.total + a
      b
    }

    /**
     * 合并缓冲区
     */
    override def merge(b1: Buff, b2: Buff): Buff = {

      b1.count = b1.count + b2.count
      b1.total = b1.total + b2.total
      b1
    }

    /**
     * 计算结果
     *
     * @param reduction
     * @return
     */
    override def finish(reduction: Buff): Long = {

      reduction.total / reduction.count
    }

    // 缓冲区编码操作
    override def bufferEncoder: Encoder[Buff] = Encoders.product

    // 输出的编码操作
    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }

}

